import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score, RocCurveDisplay,classification_report
from xgboost import XGBClassifier
from sklearn.utils import class_weight
import joblib

# 解决中文乱码问题
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

data = pd.read_csv('../../data/test2.csv')

# 删除无用列
data = data.drop(['Over18', 'StandardHours', 'EmployeeNumber'], axis=1)

data['EducationField'] = LabelEncoder().fit_transform(data['EducationField'])
data['JobRole'] = LabelEncoder().fit_transform(data['JobRole'])
data['MaritalStatus'] = LabelEncoder().fit_transform(data['MaritalStatus'])

# 采用mapping映射方法
# 手动定义映射关系
mapping1 = {'Non-Travel': 0, 'Travel_Rarely': 1, 'Travel_Frequently': 2}
# 使用map()替换，默认
data['BusinessTravel'] = data['BusinessTravel'].map(mapping1)

mapping2 = {'Human Resources': 1, 'Research & Development': 2, 'Sales': 3}
data['Department'] = data['Department'].map(mapping2)

# 热编码处理性别和是否加班列
data = pd.get_dummies(data, columns=['Gender', 'OverTime'], drop_first=True)
cols = ['Attrition'] + [col for col in data.columns if col != 'Attrition']
data = data[cols]
# 数据切割

x = data.iloc[:, 1:]
y = data.iloc[:, 0]

# class_weight = class_weight.compute_sample_weight('balanced', y_train)

model = joblib.load('../model/xgb.pkl')

y_pred = model.predict(x)
y_pre = model.predict_proba(x)[:, 1]

print(f'roc_auc_score:{roc_auc_score(y, y_pre)}')
print(f"分类评估报告:{classification_report(y, y_pred)}")

RocCurveDisplay.from_estimator(model, x, y,plot_chance_level=True)
plt.show()